#include "google/protobuf/descriptor.h"
#include "google/protobuf/message.h"
#include "base/logging.h"
#include "base/mutex.h"
#include "base/scoped_ptr.h"
#include "net/rpc/rpc_service.h"
#include "net/rpc/rpcmessage.pb.h"
#include "net/rpc/rpc.h"
#include "thread/executor.h"

using google::protobuf::Service;
using google::protobuf::RpcController;
using google::protobuf::Message;
using google::protobuf::MethodDescriptor;
using std::max;

RpcService::RpcService(int port, Executor* executor)
  : ServerConnection(port), service_(NULL), executor_(executor) {
}

void RpcService::HandleRead(int sockfd, char* buffer, int size) {
  executor_->Add(NewCallback(this, &RpcService::HandleReadInternal,
        sockfd, buffer, size));
}
void RpcService::HandleReadInternal(int sockfd, char* buffer, int size) {
  ClientContext* context = GetClientContext(sockfd);
  if (context == NULL) {
    LOG(ERROR) << "Got message after disconnection";
    return;
  }
  {
    WriteLock lock(&context->mutex);
    ByteStream* read_buf = &context->read_buf;
    read_buf->Push(buffer, size);
    while (true) {
      int header_size;
      int message_size;
      if (!read_buf->PeekTwoInt(&header_size, &message_size)) break;
      if (read_buf->size() < sizeof(int) * 2 + header_size + message_size) break;
      char nosense[sizeof(int) * 2];
      read_buf->Pop(nosense, sizeof(int) * 2);
      scoped_array<char> tmp(new char[max(header_size, message_size)]);
      scoped_ptr<RpcMessageHeader> req_header(new RpcMessageHeader);
      read_buf->Pop(tmp.get(), header_size);
      req_header->ParseFromArray(tmp.get(), header_size);
      const MethodDescriptor* method =
        service_->GetDescriptor()->FindMethodByName(req_header->name());
      if (method == NULL) {
        lock.Release();
        // TODO(): not disconnect.
        Disconnect(sockfd);
        break;
      }
      scoped_ptr<Message> req_pb(service_->GetRequestPrototype(
            method).New());
      read_buf->Pop(tmp.get(), message_size);
      scoped_ptr<Message> resp_pb (service_->GetResponsePrototype(
            method).New());
      req_pb->ParseFromArray(tmp.get(), message_size);
      VLOG(1) << req_header->DebugString();
      VLOG(1) << req_pb->DebugString();
      Message* resp_pb2 = resp_pb.release();
      Closure* closure = NewCallback(this, &RpcService::RpcDoneClosure, sockfd, req_header.release(), resp_pb2);
      RPC rpc;
      service_->CallMethod(method, &rpc, req_pb.get(),
          resp_pb2, closure);
    }
  }
  ClearDisconnected();
}

void RpcService::ClearDisconnected() {
  WriteLock lock(&client_context_mutex_);
  for (hash_map<int, ClientContext*>::iterator iter = disconnected_client_context_.begin();
      iter != disconnected_client_context_.end(); ++iter) {
    delete iter->second;
  }
  disconnected_client_context_.clear();
}

void RpcService::RpcDoneClosure(int sockfd,
    RpcMessageHeader* req_header,
    Message* resp_pb) {
  RpcMessageHeader resp_header;
  resp_header.set_id(req_header->id());
  string s1, s2;
  resp_header.SerializeToString(&s1);
  resp_pb->SerializeToString(&s2);
  int size = s1.size();
  char buf[sizeof(int)];
  memcpy(buf, (char*)&size, sizeof(int));
  Send(sockfd, buf, sizeof(int));
  size = s2.size();
  memcpy(buf, (char*)&size, sizeof(int));
  Send(sockfd, buf, sizeof(int));
  Send(sockfd, s1.c_str(), s1.size());
  Send(sockfd, s2.c_str(), s2.size());
  VLOG(1) << resp_pb->DebugString();
  delete req_header;
  delete resp_pb;
}

void RpcService::ExportService(Service* service) {
  CHECK(service_ == NULL);
  service_ = service;
}

ClientContext* RpcService::GetClientContext(int sockfd) {
  ReadLock lock(&client_context_mutex_);
  hash_map<int, ClientContext*>::const_iterator iter = client_context_.find(sockfd);
  if (iter == client_context_.end()) {
    return NULL;
  }
  return iter->second;
}

void RpcService::HandleConnect(int sockfd) {
  WriteLock lock(&client_context_mutex_);
  hash_map<int, ClientContext*>::const_iterator iter = client_context_.find(sockfd);
  if (iter == client_context_.end()) {
    client_context_[sockfd] = new ClientContext();
  }
}

void RpcService::HandleDisconnect(int sockfd) {
  VLOG(1) << "Disconnect";
  WriteLock lock(&client_context_mutex_);
  hash_map<int, ClientContext*>::const_iterator iter = client_context_.find(sockfd);
  if (iter != client_context_.end()) {
    disconnected_client_context_[iter->first] = iter->second;
    client_context_.erase(iter->first);
  }
  VLOG(1) << "Disconnected";
}

ClientContext::ClientContext() {
};


